package com.gitee.wsl.collections.iterator

import com.gitee.wsl.func.Function

class FlatCollectIterator<T, V>(
    private val iterator: Iterator<T>,
    newFunction: Function<in T, out Iterable<V>>,
) : Iterator<V> {
    private val function: Function<in T, out Iterable<V>> = newFunction
    private var innerIterator: Iterator<V> = EmptyIterator

    constructor(
        iterable: Iterable<T>,
        newFunction: Function<in T, out Iterable<V>>,
    ) : this(iterable.iterator(), newFunction)


    override fun hasNext(): Boolean {
        while (true) {
            if (this.innerIterator.hasNext()) {
                return true
            }
            if (!this.iterator.hasNext()) {
                return false
            }
            this.innerIterator = this.function(this.iterator.next()).iterator()
        }
    }

    override fun next(): V {
        if (!this.hasNext()) {
            throw NoSuchElementException()
        }
        return this.innerIterator.next()
    }
}
